# reservoir：储水池。是河、水库的类
from const import *
import numpy as np

class Reservoir(object):
    def __init__(self, serial=0, name='储水池', firstlevel=0, upconstruct_serial=None, downconstruct_serial=None,
                 reservoirtype='reservoir', levelvolumerelation=None):
        if downconstruct_serial is None:
            downconstruct_serial = []
        if upconstruct_serial is None:
            upconstruct_serial = []
        if levelvolumerelation is None:
            levelvolumerelation = [[], []]
        self.serial = serial  # 每个均应有序号
        self.name = name   # 取名
        self.reservoirtype = reservoirtype  # 默认为水库类型（另外有河流类型）
        self.upconstruct_serial = upconstruct_serial  # 上游建筑物序号
        self.downconstruct_serial = downconstruct_serial    # 下游建筑物序号
        self.waterlevellist = []  # 给出水位序列
        self.watervolumelist = []    # 给出库容序列
        self.firstlevel = firstlevel
        self.waterlevellist.append(self.firstlevel)  # 初始水位
        self.levelvolumerelation = levelvolumerelation
        if reservoirtype == 'reservoir':
            self.watervolumelist.append(self.leveltovolume(firstlevel))    # 初始库容
        else:
            self.watervolumelist = [float('inf')]    # 当为河流时，库容为无穷大

    def waterchange(self, flow):
        # 更新水位（流量计算应该放置于水利工程中）
        if self.reservoirtype == 'river':
            # 当为河流时水位保持不变
            self.waterlevellist.append(self.waterlevellist[-1])
            return 0
        else:
            volume0 = self.leveltovolume(self.waterlevellist[-1])  # 计算上时段末水位对应的库容作为初始水量
            volume1 = volume0 + flow
            # 判断水量是否为负数
            if volume1<0:
                volume1 = 0
            waterlevel1 = self.volumetolevel(volume1)
            self.waterlevellist.append(waterlevel1)
            self.watervolumelist.append(volume1)
            return volume0 + flow - volume1    # 当水量可供调蓄时，返回值应为0，当水量不足调蓄，返回差值

    def leveltovolume(self, waterlevel):
        # 根据水位计算水量
        if self.reservoirtype == 'river':
            raise Exception('正在计算河道的水量，请直接调用河道水位成果')
        else:
            # 请根据水位库容曲线计算
            levellist = np.array(self.levelvolumerelation[0])
            volumelist = np.array(self.levelvolumerelation[1])
            volume = np.interp(waterlevel, levellist, volumelist)
            if volume < 0:
                raise ValueError('水量小于零，请检查，当前水位为%s' % waterlevel)
            return volume

    def volumetolevel(self, volume):
        # 请根据水位库容曲线计算
        levellist = np.array(self.levelvolumerelation[0])
        volumelist = np.array(self.levelvolumerelation[1])
        waterlevel = np.interp(volume, volumelist, levellist)
        return waterlevel

